import warnings
warnings.filterwarnings('ignore')
from crewai import Agent, Crew, Task, LLM
import os
from dotenv import load_dotenv

# 配置API密钥
os.environ["OPENAI_API_KEY"] = "sk-a2b633c6da8c445bbe7ec80b404552b8"
load_dotenv()

# 初始化LLM
llm = LLM(
    model="deepseek/deepseek-chat",
    timeout=500,
    api_base="https://api.deepseek.com/v1",
    headers={
        "Authorization": f"Bearer {os.environ['OPENAI_API_KEY']}"
    }
)

# 定义智能体
story_outliner = Agent(
    role="故事内容构思者",
    goal="根据主题创建包含四部分的故事框架，聚焦孙悟空的成长经历",
    backstory="一位精通儿童教育的故事架构专家，擅长将成长主题转化为生动的故事框架",
    llm=llm,
    verbose=True
)

story_writer = Agent(
    role="故事内容协作者",
    goal="根据故事框架扩展完整故事，添加插画提示",
    backstory="一位擅长创作儿童文学的故事作家，能够将教育意义融入生动的情节",
    llm=llm,
    verbose=True
)

# 定义任务
outline_task = Task(
    description="""根据《西游记》主题创建故事框架：
    1. 分为四个成长阶段：拜师学艺→大闹天宫→被压五行山→西天取经
    2. 每个阶段突出一个教育主题：勇气→智慧→责任→团结
    3. 每个阶段需要包含3个关键情节点""",
    agent=story_outliner,
    expected_output="包含四个章节标题和情节要点的Markdown格式框架"
)

writing_task = Task(
    description="""将故事框架扩展为完整故事：
    1. 每段500字左右，总长度不超过3000字
    2. 每段后添加插画提示：描述场景构图、主要角色动作、色彩搭配建议
    3. 使用生动形象的语言，适合6-12岁儿童理解""",
    agent=story_writer,
    expected_output="符合格式要求的完整Markdown文档",
    context=[outline_task]
)

# 创建并运行工作流
story_crew = Crew(
    agents=[story_outliner, story_writer],
    tasks=[outline_task, writing_task],
    verbose=True
)

result = story_crew.kickoff()

# 假设 CrewOutput 对象有 content 属性，获取字符串内容
if hasattr(result, 'content'):
    result_content = result.content
else:
    # 如果没有 content 属性，尝试将对象转换为字符串
    result_content = str(result)

# 保存结果
with open("journey_to_the_west_story.md", "w", encoding="utf-8") as f:
    f.write(result_content)

print("故事生成完成！请查看 journey_to_the_west_story.md 文件")
